import pandas as pd
import pickle
from sklearn.feature_extraction.text import TfidfVectorizer
# from sklearn.model_selection import train_test_split
from tqdm import tqdm
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, classification_report
# 从sklearn.metrics中导入混淆矩阵包
from sklearn.metrics import confusion_matrix
from configs.Config import Config
import time

# 配置对象
conf = Config()
# 斌子小技巧
pd.set_option('display.expand_frame_repr', False)  # 完全多行格式显示 DataFrame,避免出现类似 ... 这样的省略号。
pd.set_option('display.max_columns', None)  # 确保所有列可见

# todo 第一步：读取数据
# 读取训练数据集
train_df_data = pd.read_csv(conf.process_train_datapath, sep='\t')
train_words = train_df_data["words"]
train_labels = train_df_data["label"]
print(f"train_df_data-->{train_df_data.head(5)}")
# print(f"train_df_data的长度为：-->{len(train_words)}")

test_df_data = pd.read_csv(conf.process_test_datapath, sep='\t')
test_words = test_df_data["words"]
test_labels = test_df_data["label"]
# print(f"test_df_data的长度为：-->{len(test_words)}")

# todo 第二步：数据预处理
# 读取停用词
# stopwords = open(conf.stop_words_path,encoding='utf-8').read().split()
stop_words = [line.strip() for line in open(conf.stop_words_path, encoding='utf-8').readlines()]
# print(stop_words)
# 创建TfidfVectorizer
tfidf = TfidfVectorizer(stop_words=stop_words)
# 将文本转为词频矩阵
train_features = tfidf.fit_transform(train_words)
test_features = tfidf.transform(test_words)
print(f"特征矩阵的维度为：{train_features.shape}")
print(train_features)  # 查看tfidf训练后特征效果
# 转化为矩阵形式展示
# # 打印特征矩阵，将稀疏矩阵转换为密集格式并输出
# print(features.toarray())
# # 打印TF-IDF模型中的所有特征名称列表
# print(list(tfidf.get_feature_names_out()))
# # 打印特征名称的数量
# print(len(tfidf.get_feature_names_out()))
# # 打印TF-IDF模型的词汇表，包含所有特征及其对应的索引
# print(tfidf.vocabulary_)
# # 打印词汇表的大小，即特征数量
# print(len(tfidf.vocabulary_))


print("*" * 100, "\n\n")

# todo 第三步：模型训练和评估
# 划分训练集和测试集
# x_train, x_test, y_train, y_test = train_test_split(features, labels, test_size=0.2, random_state=28)
x_train = train_features
y_train = train_labels
x_test = test_features
y_test = test_labels
# 训练模型
rf = RandomForestClassifier()
print("开始训练模型...")
start_time = time.time()
# 使用tqdm包装model.fit来显示进度条
for _ in tqdm(range(1), desc="随机森林模型训练进度..."):
    rf.fit(x_train, y_train)
end_time = time.time()
print(f"模型训练完成，耗时：{end_time - start_time:.2f}秒")

# 模型预测
print("开始模型预测和评估...")
y_pred = rf.predict(x_test)
# 模型评估
print("模型评估结果如下:")
print("准确率：", accuracy_score(y_test, y_pred))
print("精确率：", precision_score(y_test, y_pred, average='weighted'))
print("召回率：", recall_score(y_test, y_pred, average='weighted'))
print("F1值：", f1_score(y_test, y_pred, average='weighted'))
# # 打印评估报告
print("评估报告：", classification_report(y_test, y_pred))
# # 打印混淆矩阵
print("混淆矩阵：", confusion_matrix(y_test, y_pred))

# todo 第四步：保存模型
print("开始保存模型和向量化器...")
# 保存模型
with open(conf.random_forest_model_save_path, 'wb') as f:
    pickle.dump(rf, f)
print("模型已保存")

# 保存向量化器
with open(conf.tfidf_model_save_path, 'wb') as f:
    pickle.dump(tfidf, f)
print("向量化器已保存")
